'''
 # @ Author: Justin
 # @ Create Time: 2022-03-15 15:30:41
 # @ Modified by: Justin
 # @ Modified time: 2022-03-15 15:30:45
 # @ Description: DeepSeek实例
 '''

from openai import OpenAI
from .messages import Messages
from .messages_manager import MessagesManager
import time
import json

class DeepSeek:
    def __init__(self, 
                 CallMethod : str = 'cloud', 
                 APIkey : str = None,
                 Module : str = 'deepseek-chat',
                 url : str = 'https://api.deepseek.com',
                 frequency_penalty : float = 0.0,
                 max_tokens : int = 128,
                 presence_penalty : float = 0.0,
                 response_format : str = 'text',
                 stop : list[str] = None,
                 temperature : float = 0.9,
                 top_p : float = 1.0,
                 ):
        '''
        Args:
            CallMethod (str, optional): 调用方式，可选'cloud'或'local'. cloud表示通过API调用，local表示本地调用. 
                        默认为'cloud'.

            APIkey (str, optional): 调用API的密钥. CallMethod为'cloud'时必填.

            Module (str, optional): 调用的模块. 默认为'deepseek-chat'.

            url (str, optional): API的URL. 默认为'https://api.deepseek.com'.

            frequency_penalty (float, optional): 介于 -2.0 和 2.0 之间的数字。如果该值为正，那么新 token 会根据其
                        在已有文本中的出现频率受到相应的惩罚，降低模型重复相同内容的可能性. 默认为0.0.

            max_tokens (int, optional): 生成的文本的最大长度. 默认为128.

            presence_penalty (float, optional): 介于 -2.0 和 2.0 之间的数字。如果该值为正，那么新 token 会根据其是否
                        已在已有文本中出现受到相应的惩罚，从而增加模型谈论新主题的可能性. 默认为0.0.

            response_format (str, optional): 输出的格式，可选'text'或'json'. 'text'表示输出的文本，'json'表示输出的
                        Json. 默认为'text'.
            
            stop (list, optional):一个 string 或最多包含 16 个 string 的 list，在遇到这些词时，API 将停止生成更多的 
                        tokens. 默认为None.

            temperature (float, optional): 采样温度，介于 0 和 2 之间。更高的值，如 0.8，会使输出更随机，而更低的值，
                                    如 0.2，会使其更加集中和确定。. 默认为0.9.

            top_p (float, optional): 控制生成文本的多样性. 默认为1.0.
            
            
        '''
        
        self.CallMethod = CallMethod
        self.APIkey = APIkey
        self.Module = Module
        self.url = url
        self.frequency_penalty = frequency_penalty
        self.max_tokens = max_tokens
        self.presence_penalty = presence_penalty
        self.response_format = response_format
        self.stop = stop
        self.temperature = temperature
        self.top_p = top_p

        # 初始化OpenAI客户端
        self.client = OpenAI(
            base_url=self.url,
            api_key=self.APIkey
        )
    
    @staticmethod
    def get_cloud_deepseek(APIkey : str, Module : str = 'deepseek-chat', url : str = 'https://api.deepseek.com'):
        '''
        获取DeepSeek云端的实例
        Args:
            APIkey (str): 调用API的密钥.
            Module (str, optional): 调用的模块. 默认为'deepseek-chat'若需要r1模型则为'deepseek-reasoner'.
            url (str, optional): API的URL. 默认为DeepSeek官网的URL.'https://api.deepseek.com'.
        Returns:
            DeepSeek: DeepSeek云端的实例
        '''

        return DeepSeek(
            CallMethod='cloud', 
            APIkey=APIkey, 
            Module=Module, 
            url=url
        )
    
    @staticmethod
    def get_local_deepseek(module : str = 'deepseek-r1:1.5b', url : str = 'http://localhost:11434/v1/'):
        '''
        获取DeepSeek本地的实例
        Args:
            module (str): 调用的模块，根据本地部署的版本选择，默认'deepseek-r1:1.5b'.
            url (str, optional): API的URL. 默认为'http://localhost:11434/v1/'.
        Returns:
            DeepSeek: DeepSeek本地的实例
        '''

        return DeepSeek(
            CallMethod='local', 
            APIkey='ollama', 
            Module=module, 
            url=url
        )
    
    def chat(self, prompt : str) -> str:
        '''
        向DeepSeek发送一条消息，并获取回复
        Args:
            prompt (str): 输入的消息
        Returns:
            str: DeepSeek的回复
        '''

        try:
            response = self.client.chat.completions.create(
                model=self.Module,
                messages=[
                    {
                        'role':'user',
                        'content':prompt
                    }
                ],
                frequency_penalty=self.frequency_penalty,
                max_tokens=self.max_tokens,
                presence_penalty=self.presence_penalty,
                stop=self.stop,
                temperature=self.temperature,
                top_p=self.top_p
            )
            content = response.choices[0].message.content
            return content
        except Exception as e:
            print(f"Error: {e}")
            return "抱歉，我遇到了一个错误，请稍后再试。"
        
    def chat(self, messages : Messages, prompt : str) -> str:
        '''
        使用聊天记录与DeepSeek对话
        Args:
            messages (Messages): 聊天记录
            prompt (str): 输入的消息
        Returns:
            str: DeepSeek的回复
        '''

        try:
            messages.add_message(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), 'user', prompt)

            response = self.client.chat.completions.create(
                model=self.Module,
                messages=messages.get_chat_messages(),
                frequency_penalty=self.frequency_penalty,
                max_tokens=self.max_tokens,
                presence_penalty=self.presence_penalty,
                stop=self.stop,
                temperature=self.temperature,
                top_p=self.top_p
            )
            content = response.choices[0].message.content
            c_time = response.created

            messages.add_message(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), 'assistant', content)
            return content
        except Exception as e:
            print(f"Error: {e}")
            return "抱歉，我遇到了一个错误，请稍后再试。"
        
    def chat_with_json(self, messages : Messages, prompt : str) -> json:
        '''
        使用聊天记录与DeepSeek对话，并返回Json格式的结果
        Args:
            messages (Messages): 聊天记录
            prompt (str): 输入的消息
        Returns:
            json: DeepSeek的Json格式的回复
        '''

        # 若response_format不是json，报错
        if self.response_format != 'json':
            TypeError("response_format must be 'json'!")
        
        # 若没有prompt，报错
        if messages.prompt is None:
            ValueError("No prompt in messages can not use chat_with_json!")

        try:
            messages.add_message(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), 'user', prompt)

            response = self.client.chat.completions.create(
                model=self.Module,
                messages=messages.get_chat_messages(),
                frequency_penalty=self.frequency_penalty,
                max_tokens=self.max_tokens,
                presence_penalty=self.presence_penalty,
                stop=self.stop,
                temperature=self.temperature,
                top_p=self.top_p,
                response_format={
                    'type': 'json_object'
                }
            )
            content = response.choices[0].message.content
            c_time = response.created

            messages.add_message(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), 'assistant', content)
            return json.loads(content)
        except Exception as e:
            print(f"Error: {e}")
            return {'content':"抱歉，我遇到了一个错误，请稍后再试。", 'created':time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())}

        
    def set_frequency_penalty(self, frequency_penalty : float) -> None:
        '''
        设置频率惩罚
        Args:
            frequency_penalty (float): 介于 -2.0 和 2.0 之间的数字。如果该值为正，那么新 token 会根据其在已有文本中的出现频率受到相应的惩罚，降低模型重复相同内容的可能性.
        '''

        self.frequency_penalty = frequency_penalty
        
    def set_max_tokens(self, max_tokens : int) -> None:
        '''
        设置生成文本的最大长度
        Args:
            max_tokens (int): 生成的文本的最大长度
        '''

        self.max_tokens = max_tokens
        
    def set_presence_penalty(self, presence_penalty : float) -> None:
        '''
        设置存在惩罚
        Args:
            presence_penalty (float): 介于 -2.0 和 2.0 之间的数字。如果该值为正，那么新 token 会根据其是否已在已有文本中出现受到相应的惩罚，从而增加模型谈论新主题的可能性.
        '''

        self.presence_penalty = presence_penalty
        
    def set_response_format(self, response_format : str) -> None:
        '''
        设置输出的格式
        Args:
            response_format (str): 输出的格式，可选'text'或'json'. 'text'表示输出的文本，'json'表示输出的Json.
        '''

        self.response_format = response_format
        
    def set_stop(self, stop : list[str]) -> None:
        '''
        设置停止词
        Args:
            stop (list): 一个 string 或最多包含 16 个 string 的 list，在遇到这些词时，API 将停止生成更多的 tokens.
        '''

        self.stop = stop
        
    def set_temperature(self, temperature : float) -> None:
        '''
        设置生成文本的多样性
        Args:
            temperature (float): 控制生成文本的多样性.
        '''

        self.temperature = temperature
        
    def set_top_p(self, top_p : float) -> None:
        '''
        设置生成文本的多样性
        Args:
            top_p (float): 控制生成文本的多样性.
        '''
        self.top_p = top_p